import requests,yt_dlp
import xml.etree.ElementTree as ET
from tooles.coze.worker import Worker as Worker_coze

class Worker(object):
    def __init__(self):
        self.yt_url = "123"
        self.worker_coze = Worker_coze()

    def get_contents(self, url):
        "从srv2获取内容,[(1,2,'text')] 数组"
        response = requests.get(url)
        root = ET.fromstring(response.content)

        contents = []
        for text_element in root.findall('text'):
            text = text_element.text.strip()
            if text:
                contents.append((text_element.get('t'), text_element.get('d'), text))
        return contents
        # contnet  = " ".join([item[2] for item in contents])

    def get_vtt(self, contents):
        # 获取带标点的内容
        content = " ".join([item[2] for item in contents])
        new_content  =self.worker_coze.proc("chat2text",f"请给这段话加上合适的标点：{content}")
        return new_content

    def proc(self, yt_url):
        result = ""
        self.yt_url = yt_url
        try:
            ydl_opts = {
                'writesubtitles': True
            }
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(self.yt_url, download=False)
                if info.get("subtitles") and info.get("subtitles").get("en"):
                    subtitles = info.get("subtitles").get("en")
                else:
                    subtitles = info.get("automatic_captions").get("en")
                for subtitle in subtitles:
                    print(subtitle)
                    if subtitle.get("ext") == "srv2":
                        result = self.get_vtt(self.get_contents(subtitle["url"]))
        except Exception as e:
            print(e)
        return result

if __name__ == '__main__':
    worker = Worker()
    res = worker.proc("https://www.youtube.com/watch?v=GOMN8Ywh3c8")
    print(res)
